package com.afdor.rws.core.id.generator.impl.snowflake.zookeeper;

import com.afdor.rws.annotation.extension.Extension;
import com.afdor.rws.core.id.generator.impl.snowflake.SnowflakeIdGenerator;
import com.afdor.rws.core.strategy.IdGeneratorStrategy;
import lombok.EqualsAndHashCode;
import lombok.NoArgsConstructor;
import lombok.ToString;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.math.NumberUtils;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * 推特雪花片算法 + Zookeeper实现
 * 一个long类型的数据，64位。以下是每位的具体含义。
 * <br/>
 * snowflake的结构如下(每部分用-分开):
 * <br/>
 * 0 - 0000000000 0000000000 0000000000 0000000000 0 - 00000 - 00000 - 000000000000
 * <br/>
 * （1）第一位为未使用
 * （2）接下来的41位为毫秒级时间(41位的长度可以使用69年)
 * （3）然后是5位datacenterId
 * （4）5位workerId
 * （5）最后12位是毫秒内的计数（12位的计数顺序号支持每个节点每毫秒产生4096个ID序号）
 * <br/>
 * 一共加起来刚好64位，为一个Long型。(转换成字符串长度为18)
 *
 * @author 悭梵
 * @date Created in 2018-09-05 14:08
 */
@ToString
@EqualsAndHashCode
@NoArgsConstructor
@Extension(value = "snowflakeZookeeperId", order = Extension.NORMAL_ORDER - 80, interfaceClass = IdGeneratorStrategy.class)
public class SnowflakeZookeeperIdGeneratorStrategyImpl extends SnowflakeIdGenerator implements IdGeneratorStrategy<Long> {


    private static final String WORKER_ID = "WorkerId/";
    private static final String DATACENTER_ID = "DatacenterId/";
    private static final String ROOT_PATH = "/RWS/IdGenerator/";

    private volatile boolean stopGenerate = false;
    private CuratorFramework client;

    public SnowflakeZookeeperIdGeneratorStrategyImpl(final CuratorFramework client) {
        init(client);
    }

    // ==============================Public Methods==========================================
    public synchronized void init(final CuratorFramework client) {
        this.client = client;

        // 1、获取数据中心Id
        setDatacenterId(client);

        // 2、获取数据中心WorkerId
        setWorkerId(client);

        // 3、启动时间同步线程，为了高效生产ID，时间每秒同步一次，可调整同步策略
        Thread syncTimestampThread = new Thread(new ZookeeperStateSyncTimestampWorker(), "sync-timestamp-and-zookeeper-state");
        syncTimestampThread.setDaemon(true);
        syncTimestampThread.start();
    }

    @Override
    public synchronized Long generate() {
        if (stopGenerate) {
            throw new IllegalArgumentException("zookeeper client is stop");
        }
        return nextId();
    }

    @Override
    public synchronized List<Long> generate(int size) {
        if (stopGenerate) {
            throw new IllegalArgumentException("zookeeper client is stop");
        }
        List<Long> idList = new ArrayList<>(size);
        for (int i = 0; i < size; i++) {
            idList.add(nextId());
        }
        return idList;
    }

    // ==============================Private Methods==========================================

    /**
     * 通过Zookeeper来管理并设置WorkerId
     *
     * @param client
     */
    private void setWorkerId(final CuratorFramework client) {
        try {
            long workerId = 0L;
            Stat stat = client.checkExists().forPath(ROOT_PATH + WORKER_ID);
            if (stat == null) {
                try {
                    client.create().creatingParentsIfNeeded().withMode(CreateMode.EPHEMERAL_SEQUENTIAL).forPath(ROOT_PATH + WORKER_ID + 0L);
                    workerId = 0L;
                } catch (Exception e) {
                    setWorkerId(client);
                }
            } else {
                List<String> childrenPathList = client.getChildren().forPath(ROOT_PATH + WORKER_ID);
                if (childrenPathList.size() > this.maxWorkerId) {
                    throw new IllegalArgumentException(String.format("the number of worker id has reached a maximum of %d, and no more worker id can be created, but they can continue to be created in other datacenter id", this.maxWorkerId));
                }
                for (long i = 0L, size = Math.min(maxWorkerId, childrenPathList.size()); i <= size; i++) {
                    stat = client.checkExists().forPath(ROOT_PATH + WORKER_ID + i);
                    if (stat == null) {
                        try {
                            client.create().creatingParentsIfNeeded().withMode(CreateMode.EPHEMERAL_SEQUENTIAL).forPath(ROOT_PATH + WORKER_ID + i);
                            workerId = i;
                            break;
                        } catch (Exception e) {
                            setWorkerId(client);
                        }
                    }
                }
            }
            if (workerId > this.maxWorkerId || workerId < 0) {
                throw new IllegalArgumentException(String.format("worker id can't be greater than %d or less than 0", this.maxWorkerId));
            }
            this.workerId = workerId;
        } catch (Exception e) {
            if (e instanceof IllegalArgumentException) {
                throw (IllegalArgumentException) e;
            }
            throw new IllegalArgumentException("Failed to get worker id from zookeeper", e);
        }
    }

    /**
     * 从Zookeeper上获取datacenterId，若获取不到，默认为0
     *
     * @param client
     */
    private void setDatacenterId(final CuratorFramework client) {
        try {
            Stat stat = client.checkExists().forPath(ROOT_PATH + DATACENTER_ID);
            if (stat == null) {
                String datacenterId = System.getProperty("datacenterId");
                if (StringUtils.isBlank(datacenterId)) {
                    datacenterId = System.getenv("datacenterId");
                }
                if (StringUtils.isBlank(datacenterId)) {
                    datacenterId = "0";
                }
                try {
                    client.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(ROOT_PATH + DATACENTER_ID, datacenterId.getBytes());
                } catch (Exception e) {
                    setDatacenterId(client);
                }
            }

            byte[] datacenterIdBytes = client.getData().watched().forPath(ROOT_PATH + DATACENTER_ID);
            if (datacenterIdBytes != null && datacenterIdBytes.length > 0) {
                this.datacenterId = NumberUtils.createLong(new String(datacenterIdBytes));
            }
            if (this.datacenterId > this.maxDatacenterId || this.datacenterId < 0) {
                throw new IllegalArgumentException(String.format("datacenter id can't be greater than %d or less than 0", this.maxDatacenterId));
            }
        } catch (Exception e) {
            if (e instanceof IllegalArgumentException) {
                throw (IllegalArgumentException) e;
            }
            throw new IllegalArgumentException("failed to get datacenter id from zookeeper", e);
        }
    }

    class ZookeeperStateSyncTimestampWorker extends SyncTimestampWorker {

        @Override
        protected void needTodo() {
            stopGenerate = Objects.equals(CuratorFrameworkState.STARTED, SnowflakeZookeeperIdGeneratorStrategyImpl.this.client.getState());
        }
    }
}
